import requests
from datetime import datetime
from app.models.database import db, TargetRange
from app.utils.logger import logger

class TargetRangeService:
    """真实授权靶场对接服务"""
    
    @staticmethod
    def test_connection(target_range_id):
        """测试与靶场的连接"""
        try:
            # 查询靶场配置
            target_range = TargetRange.query.get(target_range_id)
            if not target_range:
                raise ValueError(f"Target range with id {target_range_id} not found")
            
            # 构建认证头部
            headers = {'Content-Type': 'application/json'}
            
            # 根据认证类型设置认证信息
            if target_range.auth_type == 'api_key' and target_range.api_key:
                headers['Authorization'] = f'Bearer {target_range.api_key}'
            elif target_range.auth_type == 'basic' and target_range.username and target_range.password:
                headers['Authorization'] = f'Basic {target_range.username}:{target_range.password}'
            
            # 发送测试请求
            test_url = f"{target_range.base_url}/api/v1/ping"  # 假设靶场有ping接口
            response = requests.get(test_url, headers=headers, timeout=10)
            
            # 更新最后连接时间
            target_range.last_connected = datetime.utcnow()
            db.session.commit()
            
            return {
                'success': True,
                'status_code': response.status_code,
                'message': 'Connection successful'
            }
            
        except Exception as e:
            logger.error(f"Failed to connect to target range {target_range_id}: {str(e)}")
            return {
                'success': False,
                'message': str(e)
            }
    
    @staticmethod
    def get_targets(target_range_id):
        """从靶场获取目标列表"""
        try:
            # 查询靶场配置
            target_range = TargetRange.query.get(target_range_id)
            if not target_range:
                raise ValueError(f"Target range with id {target_range_id} not found")
            
            # 构建认证头部
            headers = {'Content-Type': 'application/json'}
            
            # 根据认证类型设置认证信息
            if target_range.auth_type == 'api_key' and target_range.api_key:
                headers['Authorization'] = f'Bearer {target_range.api_key}'
            elif target_range.auth_type == 'basic' and target_range.username and target_range.password:
                headers['Authorization'] = f'Basic {target_range.username}:{target_range.password}'
            
            # 请求目标列表
            targets_url = f"{target_range.base_url}/api/v1/targets"
            response = requests.get(targets_url, headers=headers, timeout=10)
            response.raise_for_status()
            
            # 更新最后连接时间
            target_range.last_connected = datetime.utcnow()
            db.session.commit()
            
            return response.json()
            
        except Exception as e:
            logger.error(f"Failed to get targets from target range {target_range_id}: {str(e)}")
            raise
    
    @staticmethod
    def submit_scan_results(target_range_id, scan_results):
        """向靶场提交扫描结果"""
        try:
            # 查询靶场配置
            target_range = TargetRange.query.get(target_range_id)
            if not target_range:
                raise ValueError(f"Target range with id {target_range_id} not found")
            
            # 构建认证头部
            headers = {'Content-Type': 'application/json'}
            
            # 根据认证类型设置认证信息
            if target_range.auth_type == 'api_key' and target_range.api_key:
                headers['Authorization'] = f'Bearer {target_range.api_key}'
            elif target_range.auth_type == 'basic' and target_range.username and target_range.password:
                headers['Authorization'] = f'Basic {target_range.username}:{target_range.password}'
            
            # 提交扫描结果
            results_url = f"{target_range.base_url}/api/v1/scan-results"
            response = requests.post(results_url, json=scan_results, headers=headers, timeout=30)
            response.raise_for_status()
            
            # 更新最后连接时间
            target_range.last_connected = datetime.utcnow()
            db.session.commit()
            
            return response.json()
            
        except Exception as e:
            logger.error(f"Failed to submit scan results to target range {target_range_id}: {str(e)}")
            raise
    
    @staticmethod
    def get_scan_configuration(target_range_id, target_id):
        """获取靶场中特定目标的扫描配置"""
        try:
            # 查询靶场配置
            target_range = TargetRange.query.get(target_range_id)
            if not target_range:
                raise ValueError(f"Target range with id {target_range_id} not found")
            
            # 构建认证头部
            headers = {'Content-Type': 'application/json'}
            
            # 根据认证类型设置认证信息
            if target_range.auth_type == 'api_key' and target_range.api_key:
                headers['Authorization'] = f'Bearer {target_range.api_key}'
            elif target_range.auth_type == 'basic' and target_range.username and target_range.password:
                headers['Authorization'] = f'Basic {target_range.username}:{target_range.password}'
            
            # 请求扫描配置
            config_url = f"{target_range.base_url}/api/v1/targets/{target_id}/scan-config"
            response = requests.get(config_url, headers=headers, timeout=10)
            response.raise_for_status()
            
            # 更新最后连接时间
            target_range.last_connected = datetime.utcnow()
            db.session.commit()
            
            return response.json()
            
        except Exception as e:
            logger.error(f"Failed to get scan configuration from target range {target_range_id}: {str(e)}")
            raise